from sb3_contrib import RecurrentPPO
from stable_baselines3.common.vec_env import DummyVecEnv
from process_decis_env import ProcessDecisionEnv

# 假设已加载日志数据
log_data = [
    {"node": 0, "resource_util": [0.2, 0.1, 0.3], "time_elapsed": 0},
    # ...更多日志条目...
]

# 创建环境
env = DummyVecEnv([lambda: ProcessDecisionEnv(log_data)])

# 定义LSTM策略网络架构
policy_kwargs = {
    "lstm_hidden_size": 64,  # LSTM隐藏层维度
    "n_lstm_layers": 1,      # LSTM层数
    "shared_lstm": True,     # 是否共享Actor-Critic的LSTM
    "enable_critic": True    # 是否启用Critic网络
}

# 初始化RecurrentPPO模型
model = RecurrentPPO(
    "MlpLstmPolicy",
    env,
    policy_kwargs=policy_kwargs,
    verbose=1,
    n_steps=128,            # 每次更新的步数
    batch_size=32,          # 批大小（需能被n_steps整除）
    n_epochs=10,            # 每次更新的epoch数
    gamma=0.99,             # 折扣因子
    learning_rate=3e-4,
    clip_range=0.2          # PPO裁剪参数
)

# 训练模型
model.learn(total_timesteps=100000)

# 保存模型
model.save("models/rppo_process_optimizer")